import math
import os
import cv2
import numpy as np
import json

from src.smooth import list_smooth
from src.utils import _args_checker
from src.video import generate_video

count = 0

# path = "./output2"
# outpath = "" + path + "_out/"
body_connect = [(0, 1), (1, 2), (2, 3), (3, 4), (1, 5), (5, 6), (6, 7), (1, 8), (8, 9), (8, 12), (9, 10), (12, 13),
                (10, 11), (13, 14), (11, 24), (22, 23), (22, 24), (14, 21), (19, 14), (19, 20), (23, 24), (21, 20)]
hand_connect = [(0, 17),
                (18, 17),
                (18, 19),
                (19, 20),
                (0, 13),
                (14, 13),
                (14, 15),
                (15, 16),
                (0, 9),
                (9, 10),
                (10, 11),
                (11, 12),
                (0, 5),
                (6, 7),
                (5, 6),
                (7, 8),
                (0, 1),
                (1, 2),
                (3, 4),
                (2, 3),
                ]
cv2.namedWindow("image")
top_pic = './img/otou.png'
p = cv2.imread(top_pic, cv2.IMREAD_UNCHANGED)
orgi_landmarks = [[490, 650], [340, 467], [620, 467], [120, 567], [820, 567]]
head = cv2.imread("./img/otou_small.png", cv2.IMREAD_UNCHANGED)


def transformation_from_points(points1, points2):
    points1 = points1.astype(np.float32)
    points2 = points2.astype(np.float32)
    c1 = np.mean(points1, axis=0)
    c2 = np.mean(points2, axis=0)
    points1 -= c1
    points2 -= c2
    s1 = np.std(points1)
    s2 = np.std(points2)
    points1 /= s1
    points2 /= s2
    U, S, Vt = np.linalg.svd(points1.T * points2)
    R = (U * Vt).T
    return np.vstack([np.hstack(((s2 / s1) * R, c2.T - (s2 / s1) * R * c1.T)), np.matrix([0., 0., 1.])])


last = []


def draw(point_list, right, left, **kwargs):
    draw_hand = _args_checker(kwargs,"draw_hand",allow_type=bool,default=False)
    draw_border = _args_checker(kwargs,"draw_border",allow_type=bool,default=False)
    w = _args_checker(kwargs,"w",allow_type=int,default=1920)
    h = _args_checker(kwargs,"h",allow_type=int,default=1080)
    body_size = _args_checker(kwargs,"body_size",allow_type=int,default=40)
    other_size = _args_checker(kwargs,"other_size",allow_type=int,default=10)
    hand_size = _args_checker(kwargs,"hand_size",allow_type=int,default=1)
    border = _args_checker(kwargs,"border",allow_type=int,default=2)
    color = _args_checker(kwargs,"color",cov_to=tuple,default=(0, 0, 0))
    border_color = _args_checker(kwargs,"border_color",cov_to=tuple,default=(255, 255, 255))
    head_stabilize = _args_checker(kwargs,"head_stabilize",allow_type=bool,default=True)
    path = _args_checker(kwargs, "source_path",allow_type=str, allow_none=False)
    outpath = _args_checker(kwargs, "output_path",allow_type=str, default=os.path.join(path,"output"))

    png_img = p
    img = np.ones((h, w, 3), np.uint8)  # 生成一个空灰度图像
    img[:, :, 0] = 0
    img[:, :, 1] = 255
    img[:, :, 2] = 0
    pls = point_list
    for i, (x, y, z) in enumerate(pls):
        if i not in [15, 16, 17, 18]:
            cv2.circle(img, (int(x), int(y)), 5, (0, 0, 0), 4)
        # print(x, y)

    hand_left = []
    hand_right = []

    if draw_hand:
        # 左手
        for i in range(0, int(len(left) / 3)):
            x = left[i * 3]
            y = left[i * 3 + 1]
            z = left[i * 3 + 2]
            hand_left.append((int(x / 2), int(y / 2), z))
        for i, (x, y, z) in enumerate(hand_left):
            if i not in [15, 16, 17, 18]:
                cv2.circle(img, (int(x), int(y)), 2, (0, 0, 0), 1)
            # print(x, y)

        # 右手
        for i in range(0, int(len(right) / 3)):
            x = right[i * 3]
            y = right[i * 3 + 1]
            z = right[i * 3 + 2]
            hand_right.append((int(x / 2), int(y / 2), z))
        for i, (x, y, z) in enumerate(hand_right):
            if i not in [15, 16, 17, 18]:
                cv2.circle(img, (int(x), int(y)), 2, (0, 0, 0), 1)
            # print(x, y)
    # 腿
    for pt1, pt2 in [(9, 10), (12, 13)]:
        if draw_border:
            cv2.line(img, (int(pls[pt1][0]), int(pls[pt1][1])), (int(pls[pt2][0]), int(pls[pt2][1])), (255, 255, 255),
                     other_size + body_size)
        cv2.line(img, (int(pls[pt1][0]), int(pls[pt1][1])), (int(pls[pt2][0]), int(pls[pt2][1])), (0, 0, 0), other_size)

    # 身体
    if draw_border:
        cv2.line(img, (int(pls[1][0]), int(pls[1][1])), (int(pls[8][0]), int(pls[8][1])), (255, 255, 255),
                 body_size + border)
    cv2.line(img, (int(pls[1][0]), int(pls[1][1])), (int(pls[8][0]), int(pls[8][1])), (0, 0, 0), body_size)
    cv2.line(img, (int(pls[0][0]), int(pls[0][1])), (int(pls[1][0]), int(pls[1][1])), (0, 0, 0), other_size)
    cv2.line(img, (int(pls[1][0]), int(pls[1][1])), (int(pls[5][0]), int(pls[5][1])), (0, 0, 0), other_size)
    cv2.line(img, (int(pls[1][0]), int(pls[1][1])), (int(pls[2][0]), int(pls[2][1])), (0, 0, 0), other_size)

    if head_stabilize:
        global last
        px, py = 0, 0
        for i in [[pls[1][0], pls[1][1] - 40, pls[1][2]], last]:
            if i[2] != 0:
                if len(last) == 0:
                    last = i
                # print(math.sqrt((last[0] - i[0]) ** 2 + (last[1] - i[1]) ** 2))
                if math.sqrt((last[0] - i[0]) ** 2 + (last[1] - i[1]) ** 2) < 15:
                    px, py = last[0], last[1]
                else:
                    px, py = i[0], i[1]
                last = i
                break
        yy1 = 0
        yy2 = head.shape[0]
        xx1 = 0
        xx2 = head.shape[1]
        x1 = int(px - head.shape[0] / 2 if px - head.shape[0] / 2 > 0 else 0)
        y1 = int(py - head.shape[1] / 2 if py - head.shape[1] / 2 > 0 else 0)
        x2 = x1 + head.shape[1]
        y2 = y1 + head.shape[0]
        alpha_png = head[yy1:yy2, xx1:xx2, 3] / 255.0
        alpha_jpg = 1 - alpha_png
        for c in range(0, 3):
            img[y1:y2, x1:x2, c] = (alpha_png * head[yy1:yy2, xx1:xx2, c]) + \
                                   (img[y1:y2, x1:x2, c] * alpha_jpg
                                    )
    else:

        tar_landmarks = [pls[0], pls[15], pls[16], pls[17], pls[18]]
        org_landmarks = []
        ttar_landmarks = []
        # print(tar_landmarks)
        for idx, t in enumerate(tar_landmarks):
            if t[2] != 0:
                org_landmarks.append(orgi_landmarks[idx])
                ttar_landmarks.append([t[0], t[1]])

        # print(org_landmarks, ttar_landmarks)
        if len(ttar_landmarks) != 0:
            pts1 = np.float32(np.matrix([[point[0], point[1]] for point in org_landmarks]))
            pts2 = np.float32(np.matrix([[point[0], point[1]] for point in ttar_landmarks]))
            M = transformation_from_points(pts1, pts2)
            png_img = cv2.warpAffine(png_img, M[:2], (png_img.shape[0], png_img.shape[1]))
            yy1 = 0
            yy2 = png_img.shape[0]
            xx1 = 0
            xx2 = png_img.shape[1]
            x1 = 0
            y1 = 0
            x2 = x1 + png_img.shape[1]
            y2 = y1 + png_img.shape[0]

            if x1 < 0:
                xx1 = -x1
                x1 = 0
            if y1 < 0:
                yy1 = - y1
                y1 = 0
            if x2 > img.shape[1]:
                xx2 = png_img.shape[1] - (x2 - img.shape[1])
                x2 = img.shape[1]
            if y2 > img.shape[0]:
                yy2 = png_img.shape[0] - (y2 - img.shape[0])

            alpha_png = png_img[yy1:yy2, xx1:xx2, 3] / 255.0
            alpha_jpg = 1 - alpha_png
            for c in range(0, 3):
                img[y1:y2, x1:x2, c] = (alpha_png * png_img[yy1:yy2, xx1:xx2, c]) + \
                                       (img[y1:y2, x1:x2, c] * alpha_jpg
                                        )

    for pt1, pt2 in body_connect:
        if pls[pt1][0] != 0 and pls[pt1][1] != 0 and pls[pt2][0] != 0 and pls[pt2][1] != 0:
            if (pt1, pt2) in [(1, 8), (0, 1), (1, 2), (1, 5), (9, 10), (12, 13)]:
                continue
            else:
                if (pt1, pt2) in [(2, 3), (3, 4), (5, 6), (6, 7), (10, 11), (13, 14)] and draw_border:
                    cv2.line(img, (int(pls[pt1][0]), int(pls[pt1][1])), (int(pls[pt2][0]), int(pls[pt2][1])),
                             border_color, other_size + body_size)
                cv2.line(img, (int(pls[pt1][0]), int(pls[pt1][1])), (int(pls[pt2][0]), int(pls[pt2][1])), color,
                         other_size)

    if draw_hand:
        if hand_right != []:
            for pt1, pt2 in hand_connect:
                if hand_right[pt1][2] != 0 and hand_right[pt2][2] != 0:
                    cv2.line(img, (hand_right[pt1][0], hand_right[pt1][1]), (hand_right[pt2][0], hand_right[pt2][1]),
                             color, hand_size)
        if hand_left != []:
            for pt1, pt2 in hand_connect:
                if hand_left[pt1][2] != 0 and hand_left[pt2][2] != 0:
                    cv2.line(img, (hand_left[pt1][0], hand_left[pt1][1]), (hand_left[pt2][0], hand_left[pt2][1]), color,
                             hand_size)
    global count
    img_save = img
    # img_save = cv2.bitwise_not(img)
    if not os.path.exists(outpath):
        os.mkdir(outpath)
    cv2.imwrite(os.path.join(outpath,f"pic_{str(count)}.png"), img_save)
    count += 1
    cv2.imshow('image', img)

    cv2.waitKey(25)


def draw_matches(path, **kwargs):
    kp = {}
    lh = []
    rh = []

    is_kalman = _args_checker(kwargs,"is_kalman",allow_type=bool,default=True)
    is_missing_fix = _args_checker(kwargs,"is_missing_fix",allow_type=bool,default=True)
    kalman_d = _args_checker(kwargs,"kalman_d",allow_type=float,default=1.0)
    missing_type = _args_checker(kwargs,"missing_type",allow_type=int,default=1)

    # print(path)
    print('开始读取文件.....')
    for filepath, dirnames, filenames in os.walk(path):
        for filename in filenames:
            with open(os.path.join(path, filename), 'r', encoding='utf8')as jsp:
                files = json.load(jsp)
                pls = []
                if len(files["people"]) > 0:
                    point_list = files["people"][0]["pose_keypoints_2d"]
                    rh.append(files["people"][0]["hand_right_keypoints_2d"])
                    lh.append(files["people"][0]["hand_left_keypoints_2d"])
                    for i in range(0, int(len(point_list) / 3)):
                        x = point_list[i * 3]
                        y = point_list[i * 3 + 1]
                        z = point_list[i * 3 + 2]
                        pls.append((int(x / 2), int(y / 2), z))
                else:
                    pls.append((0, 0, 0))
                for i, p in enumerate(pls):
                    if kp.get(str(i), None):
                        kp[str(i)].append(p)
                    else:
                        kp[str(i)] = [p]
    for k, v in kp.items():
        kp[k] = list_smooth(v,is_kalman,is_missing_fix,kalman_d,missing_type)

    if len(kp.values()) <= 0:
        return
    for i in range(min([len(l) for l in kp.values()])):
        pls = []
        for p in kp.values():
            pls.append(p[i])
        draw(pls, rh[i], lh[i], **kwargs)
    cv2.destroyAllWindows()
    generate_video(**kwargs)
    # ax = plt.axes(projection='3d')
    # xdata = []
    # ydata = []
    # zdata = []
    # for x,y,z in kp['0'][0:10]:
    #     xdata.append(x)
    #     ydata.append(y)
    #     zdata.append(z)
    # xdata=np.array(xdata)
    # ydata=np.array(ydata)
    # zdata=np.array(zdata)
    #
    # ax.scatter3D(xdata, ydata, zdata, c=zdata, cmap='Greens')
    # plt.show()
    # draw(files["people"][0]["pose_keypoints_2d"],files["people"][0]["hand_right_keypoints_2d"],files["people"][0]["hand_left_keypoints_2d"])

def run(**kwargs):
    path = _args_checker(kwargs,"source_path",allow_type=str,allow_none=False)
    draw_matches(path,**kwargs)
